import functools
import os.path
import pickle

from datetime import datetime
from dataclasses import dataclass, replace, field
from typing import Union, Any, List, Optional

import numpy as np
from numpy import ndarray

from skfem.element import Element
from skfem.mesh import Mesh
from skfem.generic_utils import deprecated


@dataclass(repr=False)
class DofsView:
    """A subset of :class:`skfem.assembly.Dofs`."""

    obj: Any = None
    nodal_ix: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    facet_ix: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    edge_ix: Union[ndarray, slice] = field(default_factory=lambda: slice(None))
    interior_ix: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    nodal_rows: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    facet_rows: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    edge_rows: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    interior_rows: Union[ndarray, slice] = field(
        default_factory=lambda: slice(None))
    doflocs: Any = None

    def __repr__(self):
        nnodal = len(np.unique(
            self.obj.nodal_dofs[self.nodal_rows][:, self.nodal_ix]
        ))
        nfacet = len(np.unique(
            self.obj.facet_dofs[self.facet_rows][:, self.facet_ix]
        ))
        nedge = len(np.unique(
            self.obj.edge_dofs[self.edge_rows][:, self.edge_ix]
        ))
        ninterior = len(np.unique(
            self.obj.interior_dofs[self.interior_rows][:, self.interior_ix]
        ))
        dofnames = np.array(self.obj.element.dofnames)
        rep = ""
        rep += "<skfem {}({}, {}) object>".format(
            type(self).__name__,
            type(self.obj.topo).__name__,
            type(self.obj.element).__name__
        )
        if nnodal > 0:
            rep += "\n  Number of nodal DOFs: {} {}".format(
                nnodal,
                dofnames[self.nodal_rows]
            )
        if nfacet > 0:
            rep += "\n  Number of facet DOFs: {} {}".format(
                nfacet,
                dofnames[self.obj.nodal_dofs.shape[0]:][self.facet_rows],
            )
        if nedge > 0:
            rep += "\n  Number of edge DOFs: {} {}".format(
                nedge,
                dofnames[(self.obj.nodal_dofs.shape[0]
                          + self.obj.facet_dofs.shape[0]):][self.edge_rows],
            )
        if ninterior > 0:
            rep += "\n  Number of interior DOFs: {} {}".format(
                ninterior,
                dofnames[(self.obj.nodal_dofs.shape[0]
                          + self.obj.facet_dofs.shape[0]
                          + self.obj.edge_dofs.shape[0]):][self.interior_rows],
            )
        return rep

    def __str__(self):
        return self.__repr__()

    def __len__(self):
        return len(self.flatten())

    def flatten(self) -> ndarray:
        """Return all DOF indices as a single array."""
        return np.unique(
            np.concatenate((
                (self.obj
                 .nodal_dofs[self.nodal_rows][:, self.nodal_ix]
                 .flatten()),
                (self.obj
                 .facet_dofs[self.facet_rows][:, self.facet_ix]
                 .flatten()),
                (self.obj
                 .edge_dofs[self.edge_rows][:, self.edge_ix]
                 .flatten()),
                (self.obj
                 .interior_dofs[self.interior_rows][:, self.interior_ix]
                 .flatten())
            ))
        )

    def sort(self, sorting=None) -> ndarray:
        """Sort a set of DOFs based on output value of a given function.

        Parameters
        ----------
        sorting
            Optionally, specify a function which takes in a set of DOF
            locations and outputs a number.  Used for sorting the
            DOFs.  By default, the function "sum(x)" is used.

        """
        if sorting is None:
            def sorting(x):
                return sum(x)
        dofs = self.flatten()
        ix = np.argsort(sorting(self.doflocs[:, dofs]))
        return dofs[ix]

    def _intersect(self, a, b):
        if isinstance(a, slice):
            if a.start == 0:
                return a
            return b
        if isinstance(b, slice):
            if b.start == 0:
                return b
            return a
        return np.intersect1d(a, b)

    def _intersect_tuples(self, a, b):
        return tuple([self._intersect(a[i], b[i]) for i in range(len(a))])

    def keep(self, dofnames: List[str]):
        """Keep DOFs with the given names and remove others.

        Parameters
        ----------
        dofnames
            An array of DOF names, e.g. `["u", "u_n"]`.

        """
        nrows = self._intersect_tuples(
            (
                self.nodal_rows,
                self.facet_rows,
                self.edge_rows,
                self.interior_rows,
            ),
            self._dofnames_to_rows(dofnames)
        )
        return replace(
            self,
            nodal_rows=nrows[0],
            facet_rows=nrows[1],
            edge_rows=nrows[2],
            interior_rows=nrows[3],
        )

    def drop(self, dofnames):
        """Remove DOFs with the given names.

        Parameters
        ----------
        dofnames
            An array of DOF names, e.g. `["u", "u_n"]`.

        """
        nrows = self._intersect_tuples(
            (
                self.nodal_rows,
                self.facet_rows,
                self.edge_rows,
                self.interior_rows,
            ),
            self._dofnames_to_rows(dofnames, skip=True)
        )
        return replace(
            self,
            nodal_rows=nrows[0],
            facet_rows=nrows[1],
            edge_rows=nrows[2],
            interior_rows=nrows[3],
        )

    def all(self, key=None):
        if key is None:
            return self.flatten()
        return self.keep(key).flatten()

    def __array__(self, dtype=None, copy=None):
        if dtype is not None:
            return self.flatten().astype(dtype)
        return self.flatten()

    @property
    def nodal(self):
        return self._by_name(self.nodal_dofs[self.nodal_rows],
                             ix=self.nodal_ix,
                             rows=self.nodal_rows)

    @property
    def facet(self):
        return self._by_name(self.facet_dofs[self.facet_rows],
                             off=self.nodal_dofs.shape[0],
                             ix=self.facet_ix,
                             rows=self.facet_rows)

    @property
    def edge(self):
        return self._by_name(self.edge_dofs[self.edge_rows],
                             off=(self.nodal_dofs.shape[0]
                                  + self.facet_dofs.shape[0]),
                             ix=self.edge_ix,
                             rows=self.edge_rows)

    @property
    def interior(self):
        return self._by_name(self.interior_dofs[self.interior_rows],
                             off=(self.nodal_dofs.shape[0]
                                  + self.facet_dofs.shape[0]
                                  + self.edge_dofs.shape[0]),
                             ix=self.interior_ix,
                             rows=self.interior_rows)

    def __getattr__(self, attr):
        return getattr(self.obj, attr)

    @deprecated("numpy.hstack")
    def __or__(self, other):
        return replace(
            self,
            nodal_ix=np.union1d(self.nodal_ix, other.nodal_ix),
            facet_ix=np.union1d(self.facet_ix, other.facet_ix),
            edge_ix=np.union1d(self.edge_ix, other.edge_ix),
            interior_ix=np.union1d(self.interior_ix, other.interior_ix),
        )

    def __add__(self, other):
        return self.__or__(other)


class Dofs:
    """An object containing a set of degree-of-freedom indices."""

    nodal_dofs: Optional[ndarray] = None
    facet_dofs: Optional[ndarray] = None
    edge_dofs: Optional[ndarray] = None
    interior_dofs: Optional[ndarray] = None

    element_dofs: ndarray
    N: int = 0

    topo: Mesh
    element: Element

    def __init__(self, topo, element, offset=0):

        self.topo = topo
        self.element = element

        self.nodal_dofs = np.reshape(
            np.arange(element.nodal_dofs * topo.nvertices, dtype=np.int32),
            (element.nodal_dofs, topo.nvertices),
            order='F') + offset
        offset += element.nodal_dofs * topo.nvertices

        # edge dofs
        if element.dim == 3 and element.edge_dofs > 0:
            self.edge_dofs = np.reshape(
                np.arange(element.edge_dofs * topo.nedges,
                          dtype=np.int32),
                (element.edge_dofs, topo.nedges),
                order='F') + offset
            offset += element.edge_dofs * topo.nedges
        else:
            self.edge_dofs = np.empty((0, 0), dtype=np.int32)

        # facet dofs
        if element.facet_dofs > 0:
            self.facet_dofs = np.reshape(
                np.arange(element.facet_dofs * topo.nfacets,
                          dtype=np.int32),
                (element.facet_dofs, topo.nfacets),
                order='F') + offset
            offset += element.facet_dofs * topo.nfacets
        else:
            self.facet_dofs = np.empty((0, 0), dtype=np.int32)

        # interior dofs
        self.interior_dofs = np.reshape(
            np.arange(element.interior_dofs * topo.nelements, dtype=np.int32),
            (element.interior_dofs, topo.nelements),
            order='F') + offset

        # global numbering
        self.element_dofs = np.zeros((0, topo.nelements), dtype=np.int32)

        # nodal dofs
        for itr in range(topo.t.shape[0]):
            self.element_dofs = np.vstack((
                self.element_dofs,
                self.nodal_dofs[:, topo.t[itr]]
            ))

        # edge dofs
        if element.dim == 3 and element.edge_dofs > 0:
            for itr in range(topo.t2e.shape[0]):
                self.element_dofs = np.vstack((
                    self.element_dofs,
                    self.edge_dofs[:, topo.t2e[itr]]
                ))

        # facet dofs
        if element.dim >= 2 and element.facet_dofs > 0:
            for itr in range(topo.t2f.shape[0]):
                self.element_dofs = np.vstack((
                    self.element_dofs,
                    self.facet_dofs[:, topo.t2f[itr]]
                ))

        # interior dofs
        self.element_dofs = np.vstack((self.element_dofs,
                                       self.interior_dofs))

        # total dofs
        self.N = np.max(self.element_dofs) + 1

    @property
    def _iset(self):

        import petsc4py.PETSc as petsc

        return petsc.IS().createBlock(
            bsize=1,
            indices=self._globnums,
            comm=self._comm,
        )

    @property
    def _lgmap(self):

        import petsc4py.PETSc as petsc

        return petsc.LGMap().create(
            self._globnums,
            bsize=1,
            comm=self._comm,
        )

    @classmethod
    def decompose(cls,
                  comm,
                  cache=None,
                  nparts=None):
        """Decorator for building distributed meshes.

        The wrapped function must return `(mesh, dofs)`,
        i.e. :class:`skfem.mesh.Mesh` and :class:`skfem.assembly.Dofs`
        objects.

        Domain decomposition is then automatically performed using
        METIS. The parameters to the wrapper are as follows:

        Parameters
        ----------
        comm
            The MPI communicator object, e.g.,
            `petsc4py.PETSc.COMM_WORLD`
        cache
            Optional filename for caching the
            decomposition. Must contain `{}`
            which is then replaced by the index
            of the subdomain.
        nparts
            Optional number of subdomains.  If given, only METIS is
            run, the decomposition is saved to files, and the process
            terminated.  This is used, e.g., for running METIS in a
            separate high-memory SLURM task.

        """

        use_cache = False
        load_cache = False
        if isinstance(cache, str):
            use_cache = True
            # detect if cache has been filled
            # by checking the existence of rank 0 submesh
            # NOTE: do not cache if 'nparts' is given
            if os.path.isfile(cache.format(0) + '.m'):
                if nparts is None:
                    load_cache = True

        def _recv(comm):

            if use_cache:
                fname = cache.format(comm.rank)

                with open(fname + '.m', "rb") as handle:
                    m = pickle.load(handle)
                with open(fname + '.dofs', "rb") as handle:
                    dofs = pickle.load(handle)
                dofs._comm = comm
            else:
                mpicomm = comm.tompi4py()
                dofs = mpicomm.recv(source=0)
                m = dofs.topo
                dofs._comm = comm

            return m, dofs

        def _decorate(og_builder):

            @functools.wraps(og_builder)
            def wrapped_builder(*args, **kwargs):

                if comm.rank == 0 and not load_cache:

                    mesh, dofs = og_builder(*args, **kwargs)
                    retval = dofs._decompose(comm,
                                             cache=cache,
                                             use_cache=use_cache,
                                             nparts=nparts)

                if use_cache:

                    comm.barrier()

                if comm.rank > 0 or load_cache:

                    retval = _recv(comm)

                return retval

            return wrapped_builder

        return _decorate

    def _decompose(self,
                   comm,
                   cache=None,
                   use_cache=None,
                   nparts=None):

        from pymetis import part_mesh, GType

        preprocess_only = False
        if nparts is None:
            nparts = comm.size
        else:
            preprocess_only = True

        print(('{}: Calling METIS to decompose into {} subdomains...'
               .format(datetime.now(), nparts)))
        _, mship, _ = part_mesh(
                nparts,
                self.topo.t.T,
                gtype=GType.DUAL,
                ncommon=len(self.topo.elem.refdom.facets[0]),
        )
        mship = np.array(mship, dtype=np.int32)
        mpicomm = comm.tompi4py()
        retval = None

        for rank in range(nparts):
            print('{}: Processing subdomain {}...'.format(datetime.now(),
                                                          rank))
            subix = np.nonzero(mship == rank)[0]
            subm = self.topo.restrict(subix)
            subdofs = Dofs(subm, self.element)
            globnums = np.zeros(subdofs.N, dtype=np.int32)
            globnums[subdofs.element_dofs] = self.element_dofs[:, subix]
            subdofs._globnums = globnums
            subdofs._nglob = self.N

            if use_cache:
                fname = cache.format(rank)
                print(('{}: Saving subdomain {} to file...'
                       .format(datetime.now(),
                               rank)))
                with open(fname + '.m', "wb") as handle:
                    pickle.dump(subm, handle)
                with open(fname + '.dofs', "wb") as handle:
                    pickle.dump(subdofs, handle)

                if rank == 0:
                    subdofs._comm = comm
                    retval = (subm, subdofs)

            else:
                if rank > 0:
                    print(('{}: Sending subdomain {} over MPI...'
                           .format(datetime.now(),
                                   rank)))
                    mpicomm.send(subdofs, rank)
                else:
                    subdofs._comm = comm
                    retval = (subm, subdofs)

        if preprocess_only:
            raise SystemExit("'nparts' has been set: decomposition "
                             "was saved to files and process terminated")

        return retval

    def l2g(self, ix: ndarray):
        """Map DOFs from local-to-global indexing.

        The output can be passed to PETSc functions/methods that
        require global indexing (almost all of them).

        Parameters
        ----------
        ix
            Array of local indices.

        """
        ix = np.array(ix, dtype=np.int32)
        return self._lgmap.apply(ix)

    def loc(self, x):
        """Local array corresponding to a distributed PETSc vector.

        Parameters
        ----------
        x
            Distributed PETSc vector.

        """
        return x.getSubVector(self._iset).getArray()

    def get_vertex_dofs(
            self,
            nodes: ndarray,
            skip_dofnames: Optional[List[str]] = None,
            doflocs: Optional[ndarray] = None
    ) -> DofsView:
        """Return a subset of DOFs corresponding to the given nodes.

        Parameters
        ----------
        nodes
            An array of vertex indices.
        skip_dofnames
            An array of dofnames to skip.

        """
        if skip_dofnames is None:
            skip_dofnames = []

        r1, r2, r3, r4 = self._dofnames_to_rows(skip_dofnames, skip=True)

        return DofsView(
            self,
            nodes,
            np.empty((0,), dtype=np.int32),
            np.empty((0,), dtype=np.int32),
            np.empty((0,), dtype=np.int32),
            r1,
            r2,
            r3,
            r4,
            doflocs
        )

    def get_element_dofs(
            self,
            elements: ndarray,
            skip_dofnames: Optional[List[str]] = None,
            doflocs: Optional[ndarray] = None
    ) -> DofsView:
        """Return a subset of DOFs corresponding to the given elements.

        Parameters
        ----------
        elements
            An array of element indices.
        skip_dofnames
            An array of dofnames to skip.

        """
        nodal_ix = (np.empty((0,), dtype=np.int32)
                    if self.element.nodal_dofs == 0
                    else np.unique(self.topo.t[:, elements]))
        edge_ix = (np.empty((0,), dtype=np.int32)
                   if self.element.edge_dofs == 0
                   else np.unique(self.topo.t2e[:, elements]))
        facet_ix = (np.empty((0,), dtype=np.int32)
                    if self.element.facet_dofs == 0
                    else np.unique(self.topo.t2f[:, elements]))
        interior_ix = elements

        if skip_dofnames is None:
            skip_dofnames = []

        r1, r2, r3, r4 = self._dofnames_to_rows(skip_dofnames, skip=True)

        return DofsView(
            self,
            nodal_ix,
            facet_ix,
            edge_ix,
            interior_ix,
            r1,
            r2,
            r3,
            r4,
            doflocs
        )

    def get_facet_dofs(
            self,
            facets: ndarray,
            skip_dofnames: Optional[List[str]] = None,
            doflocs: Optional[ndarray] = None
    ) -> DofsView:
        """Return a subset of DOFs corresponding to the given facets.

        Parameters
        ----------
        facets
            An array of facet indices.
        skip_dofnames
            An array of dofnames to skip.

        """
        if self.element.nodal_dofs > 0 or self.element.edge_dofs > 0:
            nodal_ix, edge_ix = self.topo._expand_facets(facets)

        nodal_ix = (np.empty((0,), dtype=np.int32)
                    if self.element.nodal_dofs == 0
                    else nodal_ix)
        edge_ix = (np.empty((0,), dtype=np.int32)
                   if self.element.edge_dofs == 0
                   else edge_ix)
        facet_ix = (np.empty((0,), dtype=np.int32)
                    if self.element.facet_dofs == 0
                    else facets)

        if skip_dofnames is None:
            skip_dofnames = []

        r1, r2, r3, r4 = self._dofnames_to_rows(skip_dofnames, skip=True)

        return DofsView(
            self,
            nodal_ix,
            facet_ix,
            edge_ix,
            np.empty((0,), dtype=np.int32),
            r1,
            r2,
            r3,
            r4,
            doflocs
        )

    def _by_name(self,
                 dofs: ndarray,
                 off: int = 0,
                 ix: Optional[ndarray] = None,
                 rows: Optional[Union[List[int], ndarray]] = None):

        n_dofs = dofs.shape[0]
        n_ents = dofs.shape[1] if ix is None else len(ix)

        if rows is None:
            rows = list(range(n_dofs))

        ents = {
            self.element.dofnames[rows[i] + off]: np.zeros((0, n_ents),
                                                           dtype=np.int32)
            for i in range(n_dofs)
        }
        for i in range(n_dofs):
            new_row = dofs[i] if ix is None else dofs[i, ix]
            ents[self.element.dofnames[rows[i] + off]] =\
                np.vstack((ents[self.element.dofnames[rows[i] + off]],
                           new_row))

        return {k: ents[k].flatten() for k in ents}

    def _dofnames_to_rows(self, dofnames, skip=False):

        if isinstance(dofnames, str):
            dofnames = [dofnames]

        if skip is True:
            def check(x, y):
                return x not in y
        else:
            def check(x, y):
                return x in y

        n_nodal = self.nodal_dofs.shape[0]
        n_facet = self.facet_dofs.shape[0]
        n_edge = self.edge_dofs.shape[0]
        n_interior = self.interior_dofs.shape[0]

        nodal_rows = []
        for i in range(n_nodal):
            if check(self.element.dofnames[i], dofnames):
                nodal_rows.append(i)

        facet_rows = []
        for i in range(n_facet):
            if check(self.element.dofnames[i + n_nodal], dofnames):
                facet_rows.append(i)

        edge_rows = []
        for i in range(n_edge):
            if check(self.element.dofnames[i + n_nodal + n_facet], dofnames):
                edge_rows.append(i)

        interior_rows = []
        for i in range(n_interior):
            if check(self.element.dofnames[i + n_nodal + n_facet + n_edge],
                     dofnames):
                interior_rows.append(i)

        return (
            np.array(nodal_rows) if len(nodal_rows) > 0 else slice(0, 0),
            np.array(facet_rows) if len(facet_rows) > 0 else slice(0, 0),
            np.array(edge_rows) if len(edge_rows) > 0 else slice(0, 0),
            np.array(interior_rows) if len(interior_rows) > 0 else slice(0, 0)
        )
